from gurobipy import *
import numpy as np
from traffic_engineer import *
from traffic_engineering_class import *

class QVLB(TrafficEngineer):
	## Takes in two new variables:
	## update_frequency - how much time in-between every traffic engineering events
	def __init__(self, aurora_network, all_interblock_paths, update_frequency, training_length, reduce_multihop=False, ignore_flows=None):
		TrafficEngineer.__init__(self, aurora_network, all_interblock_paths)
		self.flow_pairs_to_ignore = ignore_flows
		self.update_frequency = update_frequency
		self.training_length = training_length
		self.reduce_multihop = reduce_multihop
		return

	def get_filename_param(self):
		return "QVLB".format(self.update_frequency, self.training_length, self.num_k)

	def get_num_clusters_for_training(self):
		return self.num_k
	# returns a string description of this class
	def get_string(self):
		return "Quasi VLB"

	def traffic_engineering_class(self):
		return TrafficEngineeringClass.ROBUST_TE

	#def _compute_mlu(self, traffic_matrices, capacity_matrix, all_paths):
	## Currently does not involve another step that reduces hop count
	def _compute_path_weights_minimize_worstcase_mlu(self, logical_topology_adj_matrix, traffic_matrices):
		nblocks = self.aurora_network.get_num_blocks()
		num_tm = len(traffic_matrices)
		model = Model("Routing minimize MLU for all traffic matrices")
		model.setParam( 'OutputFlag', False )
		mlu = model.addVar(lb=0., ub=GRB.INFINITY, obj=0, vtype=GRB.CONTINUOUS, name="mlu")
		routing_weight_vars = {}
		link_capacity_constraints = [None] * num_tm
		for k in range(num_tm):
			link_capacity_constraints[k] = [None] * nblocks
			for i in range(nblocks):
				link_capacity_constraints[k][i] = [None] * nblocks
				for j in range(nblocks):
					if i != j:
						link_capacity_constraints[k][i][j] = LinExpr()
						#routing_weight_vars[path] = model.addVar(lb=0., ub=1., obj=1., vtype=GRB.CONTINUOUS, name="w_{},{}".format(i,j))
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					weight_constraint = LinExpr()
					for path in self.all_interblock_paths[i][j]:
						routing_weight_vars[path] = model.addVar(lb=0., ub=1., obj=1., vtype=GRB.CONTINUOUS, name="w_{}".format(path))
						weight_constraint += routing_weight_vars[path]
					model.addConstr(lhs=weight_constraint, sense=GRB.EQUAL, rhs=1.)


		# stage 1: setup all the optimization variables that are the routing weights
		# stage 1.5 : also adds the traffic flow satisfiability constraints
		for tm, tm_index in zip(traffic_matrices, range(num_tm)):
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						for path in self.all_interblock_paths[i][j]:
							src = path[0]
							for path_hop in range(1, len(path), 1):
								dst = path[path_hop]
								link_capacity_constraints[tm_index][src][dst] += (routing_weight_vars[path] * tm[i][j])
								src = dst
						#for k in range(nblocks):
						#	if k != i and k != j:
						#		lu_constraint += routing_weight_vars[(i,k,j)] * tm[i][k] + routing_weight_vars[(k,j,i)] * tm[k][j]
						#		weight_constraint += routing_weight_vars[(i, j, k,)]
						#model.addConstr(lhs=lu_constraint, sense=GRB.LESS_EQUAL, rhs=mlu * capacity_matrix[i][j] * logical_topology_adj_matrix[i][j])
						#model.addConstr(lhs=weight_constraint, sense=GRB.EQUAL, rhs=1.)
		## stage 2 : add the link utilization constraints to the model
		for tm_index in range(num_tm):
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						lu_constraint = link_capacity_constraints[tm_index][i][j]
						capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(j))
						model.addConstr(lhs=lu_constraint, sense=GRB.LESS_EQUAL, rhs=mlu * capacity * logical_topology_adj_matrix[i][j])
		# stage 3: set the objective
		model.setObjective(mlu, GRB.MINIMIZE)
		try: 
			model.optimize()
			routing_weights = {}
			for path in routing_weight_vars.keys():
				routing_weights[path] = routing_weight_vars[path].x
			return routing_weights, mlu.x
		except GurobiError as e:
			print ("Error code " + str(e. errno ) + ": " + str(e))
			return None
		except AttributeError :
			print ("Encountered an attribute error ")
			return None

	def minimize_burst_risk(self, logical_topology_adj_matrix, traffic_matrices, worst_case_mlu):
		routing_weights = {}
		nblocks = self.aurora_network.get_num_blocks()
		num_tm = len(traffic_matrices)
		model = Model("Robustify")
		model.setParam( 'OutputFlag', False )
		routing_weight_vars = {}
		link_capacity_constraints = [None] * num_tm
		beta = model.addVar(lb=0., ub=GRB.INFINITY, obj=1, vtype=GRB.CONTINUOUS, name="beta")

		max_tm = np.zeros((nblocks, nblocks))
		for tm in traffic_matrices:
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						max_tm[i][j] = max(tm[i][j], max_tm[i][j])

		for k in range(num_tm):
			link_capacity_constraints[k] = [None] * nblocks
			for i in range(nblocks):
				link_capacity_constraints[k][i] = [None] * nblocks
				for j in range(nblocks):
					if i != j:
						link_capacity_constraints[k][i][j] = LinExpr()
		
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					weight_constraint = LinExpr()
					for path in self.all_interblock_paths[i][j]:
						routing_weight_vars[path] = model.addVar(lb=0., ub=1., obj=0, vtype=GRB.CONTINUOUS, name="w_{}".format(path))
						weight_constraint += routing_weight_vars[path]
					model.addConstr(lhs=weight_constraint, sense=GRB.EQUAL, rhs=1.)
		# stage 1: setup all the optimization variables that are the routing weights
		# stage 1.5 : also adds the traffic flow satisfiability constraints
		for tm, tm_index in zip(traffic_matrices, range(len(traffic_matrices))):
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						for path in self.all_interblock_paths[i][j]:
							src = path[0]
							for path_hop in range(1, len(path), 1):
								dst = path[path_hop]
								link_capacity_constraints[tm_index][src][dst] += (routing_weight_vars[path] * tm[i][j])
								src = dst

		## stage 2 : add the link utilization constraints
		for tm_index in range(num_tm):
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						lu_constraint = link_capacity_constraints[tm_index][i][j]
						capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(j))
						model.addConstr(lhs=lu_constraint, sense=GRB.LESS_EQUAL, rhs=worst_case_mlu * capacity * logical_topology_adj_matrix[i][j])
						
		# stage 3: set the objective
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					model.addConstr(lhs=routing_weight_vars[(i, j)] , sense=GRB.LESS_EQUAL, rhs=beta * logical_topology_adj_matrix[i][j] * max_tm[i][j])

		model.setObjective(beta, GRB.MINIMIZE)
		try: 
			model.optimize()
			routing_weights = {}
			for path in routing_weight_vars.keys():
				routing_weights[path] = routing_weight_vars[path].x
			return routing_weights
		except GurobiError as e:
			print ("Error code " + str(e. errno ) + ": " + str(e))
			return None
		except AttributeError :
			print ("Encountered an attribute error ")
			return None
		return routing_weights

	def compute_path_weights(self, logical_topology_adj_matrix, traffic_matrices, spare_traffic=0.1, fraction_threshold=0.002):
		## newly added for robustness
		nblocks = self.aurora_network.get_num_blocks()
		egress_traffic = [0] * nblocks
		ingress_traffic = [0] * nblocks

		## Computes the largest ingress and egress entries
		for tm in traffic_matrices:
			for pod in range(nblocks):
				egress = sum(tm[pod])
				egress_traffic[pod] = max(egress, egress_traffic[pod])
				ingress = 0
				for target_pod in range(nblocks):
					ingress += tm[target_pod][pod]
				ingress_traffic[pod] = max(ingress, ingress_traffic[pod])
				
		routing_weights, worst_case_mlu = self._compute_path_weights_minimize_worstcase_mlu(logical_topology_adj_matrix, traffic_matrices)
		if self.reduce_multihop:
			#routing_weights = self._minimize_multihop_direct(logical_topology_adj_matrix, traffic_matrices, worst_case_mlu)
			routing_weights = self._minimize_multihop(logical_topology_adj_matrix, traffic_matrices, worst_case_mlu)
		#routing_weights = self.minimize_burst_risk(logical_topology_adj_matrix, traffic_matrices, worst_case_mlu)
		return routing_weights